#include "wrapper_sstructmg.h"

const idx_t dims_per_blk[] = {
     64, 640, 768,
    //  64, 160, 192,
      1,   1,   1,
      1,   1,   1,
      1,   1,   1,
      1,   1,   1,
      1,   1,   1,
};

const idx_t offd_nnz [] = {
//  0   1   2   3   4   5
    0, 64, 64, 64, 64, 64,// 0
    64, 0,  0,  0,  0,  0,// 1
    64, 0,  0,  0,  0,  0,// 2
    64, 0,  0,  0,  0,  0,// 3
    64, 0,  0,  0,  0,  0,// 4
    64, 0,  0,  0,  0,  0,// 5
};

int main(int argc, char* argv[])
{
    int my_pid, num_procs;
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
    MPI_Comm_size(MPI_COMM_WORLD, &num_procs);
    if (my_pid == 0) printf("\033[1;35mMax threads %d\033[0m \n", omp_get_max_threads());

    idx_t arg_cnt = 1;
    const idx_t pN[3] = { atoi(argv[arg_cnt++]), atoi(argv[arg_cnt++]),  atoi(argv[arg_cnt++]) };
    assert(pN[0] * pN[1] * pN[2] == num_procs);
    const std::string pathname(argv[arg_cnt++]);
    const std::string   its_name = std::string(argv[arg_cnt++]),
                        prc_name = std::string(argv[arg_cnt++]);
    std::string config_mg_file;
	if (prc_name == "GMG") config_mg_file = std::string(argv[arg_cnt++]);

    const idx_t glb_nparts = sizeof(dims_per_blk) / sizeof(idx_t) / 3;
    idx_t blk_offsets[glb_nparts + 1];
    blk_offsets[0] = 0;
    for (idx_t b = 0; b < glb_nparts; b++) {
        blk_offsets[b + 1] = blk_offsets[b] 
            + dims_per_blk[b*3  ] * dims_per_blk[b*3+1] * dims_per_blk[b*3+2];
    }
    // const idx_t tot_nrows = blk_offsets[glb_nparts];

    idx_t pCoords[3], begs[3], ends[3];
    pCoords[0] =  my_pid / (pN[1] * pN[2]);
    pCoords[1] = (my_pid - pCoords[0] * pN[1] * pN[2]) / pN[2];
    pCoords[2] =  my_pid - pCoords[0] * pN[1] * pN[2] - pCoords[1] * pN[2];

    for (idx_t d = 0; d < 3; d++) {
        idx_t loc_len = dims_per_blk[d] / pN[d];
        idx_t remain = dims_per_blk[d] - loc_len * pN[d];
        begs[d] = pCoords[d] * loc_len;
        if (pCoords[d] < remain) {
            begs[d] += pCoords[d];
            loc_len += 1;
        } else {
            begs[d] += remain;
        }
        ends[d] = begs[d] + loc_len;
    }

    std::vector<PartDescriptor<idx_t> > my_parts;
    {// 本进程负责的部分
        PartDescriptor<idx_t> descr;
        descr.id = 0;
        descr.nelems = 1;
        for (idx_t d = 0; d < 3; d++) {
            descr.ilower[d] = begs[d];
            descr.iupper[d] = ends[d] - 1;
            descr.nelems *= (descr.iupper[d] - descr.ilower[d] + 1);
        }
        my_parts.push_back(descr);

        if (begs[1] <= 319 && 319 < ends[1] &&
            begs[2] <= 383 && 383 < ends[2] && begs[0] == 0) {
        // if (begs[1] <= 79 && 79 < ends[1] &&
        //     begs[2] <= 95 && 95 < ends[2] && begs[0] == 0) {
            descr.id = 1;
            descr.nelems = 1;
            for (idx_t d = 0; d < 3; d++) { descr.ilower[d] = descr.iupper[d] = 0; }
            my_parts.push_back(descr);
        }
        if (begs[1] <=   0 &&   0 < ends[1] &&
            begs[2] <=   0 &&   0 < ends[2] && begs[0] == 0) {
            descr.id = 2;
            descr.nelems = 1;
            for (idx_t d = 0; d < 3; d++) { descr.ilower[d] = descr.iupper[d] = 0; }
            my_parts.push_back(descr);
        }
        // if (begs[1] <=  79 &&  79 < ends[1] &&
        if (begs[1] <= 319 && 319 < ends[1] &&
            begs[2] <=   0 &&   0 < ends[2] && begs[0] == 0) {
            descr.id = 3;
            descr.nelems = 1;
            for (idx_t d = 0; d < 3; d++) { descr.ilower[d] = descr.iupper[d] = 0; }
            my_parts.push_back(descr);
        }
        if (begs[1] <=   0 &&   0 < ends[1] &&
            // begs[2] <=  95 &&  95 < ends[2] && begs[0] == 0) {
            begs[2] <= 383 && 383 < ends[2] && begs[0] == 0) {
            descr.id = 4;
            descr.nelems = 1;
            for (idx_t d = 0; d < 3; d++) { descr.ilower[d] = descr.iupper[d] = 0; }
            my_parts.push_back(descr);
        }
        // if (begs[1] <= 159 && 159 < ends[1] &&
        //     begs[2] <= 191 && 191 < ends[2] && begs[0] == 0) {
        if (begs[1] <= 639 && 639 < ends[1] &&
            begs[2] <= 767 && 767 < ends[2] && begs[0] == 0) {
            descr.id = 5;
            descr.nelems = 1;
            for (idx_t d = 0; d < 3; d++) { descr.ilower[d] = descr.iupper[d] = 0; }
            my_parts.push_back(descr);
        }
    }
    SStructMG_Init();
    {
        SStructGrid<idx_t> ssgrid(MPI_COMM_WORLD, 3, glb_nparts);
        const idx_t halos[3] = { 1, 1, 1 };
        {// 添加本进程所拥有的块和盒子
            std::unordered_set<idx_t> my_added_blks;
            for (idx_t i = 0; i < (idx_t) my_parts.size(); i++) {
                const idx_t blk_id = my_parts[i].id;
                if (my_added_blks.find(blk_id) == my_added_blks.end()) {// 本进程之前没有负责过这个块
                    const idx_t blk_exts[3] = { dims_per_blk[blk_id * 3 + 0],
                                                dims_per_blk[blk_id * 3 + 1],
                                                dims_per_blk[blk_id * 3 + 2] };
                    
                    const bool periods[3] = { false, false, false };
                    // printf("Proc %d addBlock %d : (%d,%d,%d)\n", my_pid, blk_id, blk_exts[0], blk_exts[1], blk_exts[2]);
                    ssgrid.addBlock(blk_id, blk_exts, periods);
                    my_added_blks.emplace(blk_id);// 记录一下已添加该块
                }
                const idx_t blk_ends[3] = { my_parts[i].iupper[0] + 1,
                                            my_parts[i].iupper[1] + 1,
                                            my_parts[i].iupper[2] + 1 };
                // printf("Proc %d addBox to Blk %d : (%d,%d,%d) ~ (%d,%d,%d)\n",
                //     my_pid, blk_id, my_parts[i].ilower[0], my_parts[i].ilower[1], my_parts[i].ilower[2],
                //         blk_ends[0], blk_ends[1], blk_ends[2]);
                ssgrid.addBox(blk_id, my_parts[i].ilower, blk_ends);
            }
        }
        
        ssgrid.assemble();
        const idx_t my_nblks = ssgrid.my_blk_list.size();
        assert(my_nblks == (idx_t) my_parts.size());
        idx_t shalos [my_nblks * 3];
        for (idx_t k = 0; k < my_nblks; k++) {
            shalos[k * 3 + 0] = halos[0];
            shalos[k * 3 + 1] = halos[1];
            shalos[k * 3 + 2] = halos[2];
        }

        // 根据半结构网格建立向量
        par_SstructVector<idx_t, data_t> b(ssgrid, shalos, true);
        par_SstructVector<idx_t, data_t> x(b), y(b), x0(b);
        {// 读入向量数据
            for (idx_t blk = 0, ptr = 0; blk < glb_nparts; blk++) {
                if (ptr < my_nblks && ssgrid.my_blk_list[ptr].blk_id == blk) {
                    // printf("Proc %d reading blk %d\n", my_pid, blk);
                    b.struct_vectors[ptr]->IO_data(0, pathname, "b.B" + std::to_string(blk));
                    x0.struct_vectors[ptr]->IO_data(0, pathname, "x0.B"+ std::to_string(blk));
                    ptr++;
                }
                // MPI_Barrier(MPI_COMM_WORLD);
                // if (my_pid == 0) printf("Load %d done\n", blk);
            }
            y.set_val(0.0, true);
        }
        
        // 根据半结构网格建立矩阵
        __int128_t masks[my_nblks]; for (idx_t k = 0; k < my_nblks; k++) masks[k] = stencil_mask_3d7;
        par_SstructMatrix<idx_t, data_t, data_t> A(ssgrid, shalos, masks);
        {// 读入矩阵数据
            // 首先处理对角块，也即结构化部分
            for (idx_t blk = 0, ptr = 0; blk < glb_nparts; blk++) {
                if (ptr < my_nblks && ssgrid.my_blk_list[ptr].blk_id == blk) {
                    A.diag_matrixs[ptr]->IO_data(0, pathname, "B" + std::to_string(blk) + "_Diag");
                    ptr++;
                }
                // MPI_Barrier(MPI_COMM_WORLD);
                // if (my_pid == 0) printf("Load %d done\n", blk);
            }
            
            FILE * fp = nullptr;
            // 然后处理非对角快，也即非结构部分
            A.init_offd();
            
            std::vector< Entry<idx_t, data_t> > all_nnz;
            for (idx_t mk = 0; mk < my_nblks; mk++) {// 遍历每一个本进程需要参与的块
                const Block<idx_t> & blk = ssgrid.my_blk_list[mk];
                const idx_t my_nboxs = blk.my_box_list.size();
                for (idx_t nbk = 0; nbk < glb_nparts; nbk++) {// 遍历每一个可能的邻居块
                    const idx_t nnz = offd_nnz[blk.blk_id * glb_nparts + nbk];
                    if (nnz == 0) continue;
                    // 如果该邻居块对本块有影响
                    const idx_t ngb_dims[3] = { dims_per_blk[nbk * 3 + 0],
                                                dims_per_blk[nbk * 3 + 1],
                                                dims_per_blk[nbk * 3 + 2] };
                    // 读入该非对角块的数据
                    COO<idx_t, data_t> * buf = new COO<idx_t, data_t> [nnz];
                    const std::string filename = pathname + "/B" + std::to_string(blk.blk_id)+"_Offd"+std::to_string(nbk);
                    fp = fopen(filename.c_str(), "rb");
                    size_t ret = fread(buf, sizeof(COO<idx_t, data_t>), nnz, fp); assert(ret == (size_t) nnz);
                    fclose(fp);

                    StructDescriptor<idx_t> row, col;
                    row.id = blk.blk_id;
                    col.id = nbk;
                    for (idx_t k = 0; k < nnz; k++) {// 遍历每一个可能的非零元
                        const idx_t loc_1D = buf[k].r - blk_offsets[blk.blk_id],// 该非零元的行序号在本块内的一维索引
                                    ngb_1D = buf[k].c - blk_offsets[nbk];// 该非零元的列序号在邻居块内的一维索引
                        idx_t * index = row.index;// 该非零元的行在本块内的三维坐标
                        index[0] =  loc_1D / (blk.extents[1] * blk.extents[2]);
                        index[1] = (loc_1D - index[0] * blk.extents[1] * blk.extents[2]) / blk.extents[2];
                        index[2] =  loc_1D - index[0] * blk.extents[1] * blk.extents[2] - index[1] * blk.extents[2];
                        idx_t * to_id = col.index;// 该非零元的列在邻居块内的三维坐标
                        to_id[0] =  ngb_1D / (ngb_dims[1] * ngb_dims[2]);
                        to_id[1] = (ngb_1D - to_id[0] * ngb_dims[1] * ngb_dims[2]) / ngb_dims[2];
                        to_id[2] =  ngb_1D - to_id[0] * ngb_dims[1] * ngb_dims[2] - to_id[1] * ngb_dims[2];

                        for (idx_t bx = 0; bx < my_nboxs; bx++) {
                            const idx_t * begs = blk.my_box_list[bx].begs,
                                        * ends = blk.my_box_list[bx].ends;
                            if (begs[0] <= index[0] && index[0] < ends[0] &&
                                begs[1] <= index[1] && index[1] < ends[1] &&
                                begs[2] <= index[2] && index[2] < ends[2] )
                            {// 该行属于这个盒子
                                // if (row.id == 4) {
                                //     printf("Proc %d add %d (%d,%d,%d) -- %d (%d,%d,%d) %.3e\n", my_pid,
                                //         row.id, row.index[0], row.index[1], row.index[2],
                                //         col.id, col.index[0], col.index[1], col.index[2], buf[k].v);
                                // }
                                all_nnz.push_back(Entry<idx_t, data_t>(row, col, buf[k].v));
                                break;
                            }
                        }// 找到所属盒子后跳出循环
                    }
                    delete [] buf;
                }
            }
            // 是否需要混洗
            unsigned seed = std::chrono::system_clock::now().time_since_epoch().count();
            std::shuffle(all_nnz.begin(), all_nnz.end(), std::default_random_engine(seed));
            A.add_entries(all_nnz.size(), all_nnz.data());
            
        }
        A.assemble();
        
        const int test_cnt = 10;
        std::vector<TEST_RECORD> records;
        for (int test = 0; test < test_cnt; test++) {
            vec_copy(x0, x);

            double ck = vec_dot<idx_t, data_t, double>(b, b);
            if (my_pid == 0) printf("(  b,   b) = %.15e\n", ck);

            ck = vec_dot<idx_t, data_t, double>(x, x);
            if (my_pid == 0) printf("(  x,   x) = %.15e\n", ck);

            A.Mult(b, y, false);
            ck = vec_dot<idx_t, data_t, double>(y, y);
            if (my_pid == 0) printf("(A*b, A*b) = %.15e\n", ck);

            A.Mult(x, y, false);
            ck = vec_dot<idx_t, data_t, double>(y, y);
            if (my_pid == 0) printf("(A*x, A*x) = %.15e\n", ck);
            
            TEST_CONFIG config;
            config.my_nblks = my_nblks;
            config.config_mg_file = config_mg_file;
            config.ssgrid_ptr = & ssgrid;
            config.rtol = 1.0e-7;
            TEST_RECORD rec;

            buildup_solver(its_name, prc_name, config);
            setup_and_solve(A, b, x, rec);

            data_t true_r_norm, b_norm;
            check_residual(A, x, b, y, true_r_norm, b_norm);

            if (my_pid == 0) {
                printf("\033[1;35mtrue ||r|| = %20.16e ||r||/||b||= %20.16e\033[0m\n", 
                    true_r_norm, true_r_norm / b_norm);
                printf("Proc %d Setup, Solve costs %.6f %.6f s\n",
                    my_pid, rec.setup, rec.solve);
            }

            stat_part_times(rec);
            records.push_back(rec);

            destroy_solver();
        }// test loop

        // 输出最佳时间
        if (my_pid == 0) {
            stat_multi_runs(records);
        }
    }// all Semi-** object must be destroyed before MPI_Finalize()
    SStructMG_Finalize();
    MPI_Finalize();
    return 0;
}


